// Copyright  2004-2008 Backplane Ltd.
//
// This file is part of Base64
//
// Base64 is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 3 of the License, or (at your
// option) any later version.
//
// Base64 is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
// more details.
//
// You should have received a copy of the GNU General Public License
// along with Base64. If not, see <http://www.gnu.org/licenses/>.
//	
//

#include "stdafx.h"
#include <COMDEF.H>
#include <atlenc.h>
#include <string.h>
#include "Base64.h"
#include "codec.h"

// Flags for encoding options.
const char g_flags = ATL_BASE64_FLAG_NONE;

Ccodec::Ccodec()
//: m_pInput(0), m_nFileSize(0)
: m_hFile(0), m_hMap(0), m_nFileSize(0), m_lpszDecoded(0)
{
}

Ccodec::~Ccodec()
{
	closeFile();
}

STDMETHODIMP Ccodec::InterfaceSupportsErrorInfo(REFIID riid)
{
	static const IID* arr[] = 
	{
		&IID_Icodec
	};
	for (int i=0; i < sizeof(arr) / sizeof(arr[0]); i++)
	{
		if (::InlineIsEqualGUID(*arr[i],riid))
			return S_OK;
	}
	return S_FALSE;
}

STDMETHODIMP Ccodec::encodeFile(BSTR sFilename, BSTR *psRetval)
{
	HRESULT hr = E_FAIL;
	if(psRetval) 
	{
		char *lpszFileName = _com_util::ConvertBSTRToString(sFilename);
		if(lpszFileName) 
		{
			HANDLE hFile = CreateFile(lpszFileName, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, 0, 0);
			if(hFile != INVALID_HANDLE_VALUE) 
			{
				HANDLE hMap = CreateFileMapping(hFile, 0, PAGE_READONLY, 0, 0, 0);
				if(hMap) 
				{
					void *pDecoded = MapViewOfFile(hMap, FILE_MAP_READ, 0, 0, 0);
					if(pDecoded) 
					{
						DWORD nFileSize = GetFileSize(hFile, 0);
						int nWritten = Base64EncodeGetRequiredLength(nFileSize, g_flags);
						char *lpszEncoded = new char[nWritten];
						if(lpszEncoded) 
						{
							if(Base64Encode(static_cast<BYTE*>(pDecoded), nFileSize, lpszEncoded, &nWritten, g_flags)) 
							{
								CComBSTR bs(nWritten, lpszEncoded);
								*psRetval = bs.Detach(); // This memory now belongs to the caller
								hr = S_OK;
							}
							delete[] lpszEncoded;
						}
					}
					CloseHandle(hMap);
				}
				CloseHandle(hFile);
			}
			delete[] lpszFileName;
		}
	} 
	else
		hr = E_INVALIDARG;
	return hr;
}

// Optimise further.
STDMETHODIMP Ccodec::decodeText(BSTR sCoded, BSTR sFilename)
{
	HRESULT hr = E_FAIL;
	if(sCoded && sFilename) 
	{
		char *lpszFileName = _com_util::ConvertBSTRToString(sFilename);
		if(lpszFileName) 
		{
			char *lpszEncoded = _com_util::ConvertBSTRToString(sCoded);
			if(lpszEncoded) 
			{
				int nStrLen = strlen(lpszEncoded);
				char *lpszDecoded = new char[nStrLen];
				if(lpszDecoded) 
				{
					if(Base64Decode(lpszEncoded, nStrLen, reinterpret_cast<BYTE*>(lpszDecoded), &nStrLen)) 
					{
						HANDLE hFile = CreateFile(lpszFileName, GENERIC_READ | GENERIC_WRITE | DELETE, 0, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);
						if(hFile != INVALID_HANDLE_VALUE) 
						{
							HANDLE hMap = CreateFileMapping(hFile, 0, PAGE_READWRITE, 0, nStrLen, 0);
							if(hMap) 
							{
								void *pFile = MapViewOfFile(hMap, FILE_MAP_WRITE, 0, 0, 0);
								if(pFile) 
								{
									memcpy(pFile, static_cast<void*>(lpszDecoded), nStrLen);
									hr = S_OK;
								}
								CloseHandle(hMap);
							}
							CloseHandle(hFile);
						}
					}
					delete[] lpszDecoded;
				}
				delete[] lpszEncoded;
			}
			delete[] lpszFileName;
		}
	} 
	else
		hr = E_INVALIDARG;
	return hr;
}

STDMETHODIMP Ccodec::encodeText(BSTR sText, BSTR *psRetVal)
{
	HRESULT hr = E_FAIL;
	if(sText && psRetVal) 
	{
		char *lpszDecoded = _com_util::ConvertBSTRToString(sText);
		if(lpszDecoded) 
		{
			int nStrLen = strlen(lpszDecoded);
			int nWritten = Base64EncodeGetRequiredLength(nStrLen, g_flags);
			char *lpszEncoded = new char[nWritten];
			if(lpszEncoded) 
			{
				if(Base64Encode(reinterpret_cast<BYTE*>(lpszDecoded), nStrLen, lpszEncoded, &nWritten, g_flags)) 
				{
					CComBSTR bs(nWritten, lpszEncoded);
					*psRetVal = bs.Detach(); // This memory now belongs to the caller
					hr = S_OK;
				}
				delete[] lpszEncoded;
			}
			delete[] lpszDecoded;
		}
	} 
	else
		hr = E_INVALIDARG;
	return hr;
}

STDMETHODIMP Ccodec::encodeStream(BSTR *psTargetString)
{
	HRESULT hr = E_FAIL;
	if(psTargetString) 
	{
		if(m_nFileSize) 
		{
			int nWritten = Base64EncodeGetRequiredLength(m_nFileSize, g_flags);
			char *lpszEncoded = new char[nWritten];
			if(lpszEncoded) 
			{
				if(Base64Encode(reinterpret_cast<BYTE*>(m_lpszDecoded), m_nFileSize, lpszEncoded, &nWritten, g_flags)) 
				{
					CComBSTR bs(nWritten, lpszEncoded);
					*psTargetString = bs.Detach(); // This memory now belongs to the caller
					hr = S_OK;
				}
				delete[] lpszEncoded;
			}
		}
	} 
	else
		hr = E_INVALIDARG;
	return hr;
}

STDMETHODIMP Ccodec::openFile(BSTR sFilename)
{
	HRESULT hr = E_FAIL;
	if(sFilename) 
	{
		char *lpszFileName = _com_util::ConvertBSTRToString(sFilename);
		if(lpszFileName) 
		{
			if(m_hFile)
				CloseHandle(m_hFile);
			m_hFile = CreateFile(lpszFileName, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, 0, 0);
			if(m_hFile != INVALID_HANDLE_VALUE) 
			{
				if(m_hMap)
					CloseHandle(m_hMap);
				m_hMap = CreateFileMapping(m_hFile, 0, PAGE_READONLY, 0, 0, 0);
				if(m_hMap) 
				{
					m_lpszDecoded = static_cast<char*>(MapViewOfFile(m_hMap, FILE_MAP_READ, 0, 0, 0));
					if(m_lpszDecoded) 
					{
						m_nFileSize = GetFileSize(m_hFile, 0);
						hr = S_OK;
					}
				}
			}
			delete[] lpszFileName;
		}
	} 
	else
		hr = E_INVALIDARG;
	return hr;
}

STDMETHODIMP Ccodec::closeFile()
{
	HRESULT hr = S_OK;

	m_nFileSize = 0;

	if(m_hMap)
		if(CloseHandle(m_hMap))
			m_hMap = NULL;
		else
			hr = HRESULT_FROM_WIN32(GetLastError());
    if(m_hFile)
		if(CloseHandle(m_hFile))
			m_hFile = NULL;
		else 
			hr = HRESULT_FROM_WIN32(GetLastError());
	
	return hr;
}

/*
 * This was my attempt to use existing streams-based behaviour
 * with the ATL Base64 functions. I gave up because I couldn't
 * get istream::read() to work - pbooth, 20030605.
 *
STDMETHODIMP Ccodec::encodeStream(BSTR *psTargetString)
{
	HRESULT hr = E_FAIL;
	if(psTargetString) {
		if(m_pInput) {
			char *lpszDecoded = new char[m_nFileSize];
			if(lpszDecoded) {
				m_pInput->read(lpszDecoded, m_nFileSize);
				int nWritten = Base64EncodeGetRequiredLength(m_nFileSize, g_flags);
				char *lpszEncoded = new char[nWritten];
				if(lpszEncoded) {
					if(Base64Encode(reinterpret_cast<BYTE*>(lpszDecoded), m_nFileSize, lpszEncoded, &nWritten, g_flags)) {
						CComBSTR bs(nWritten, lpszEncoded);
						*psTargetString = bs.Detach(); // This memory now belongs to the caller
						hr = S_OK;
					}
					delete[] lpszEncoded;
				}
				delete[] lpszDecoded;
			}
		}
	} else
		hr = E_INVALIDARG;
	return hr;
}

STDMETHODIMP Ccodec::openFile(BSTR sFilename)
{
	HRESULT hr = E_FAIL;
	if(sFilename) {
		if(m_pInput)
			delete m_pInput;
		char *lpszFileName = _com_util::ConvertBSTRToString(sFilename);
		if(lpszFileName) {
			HANDLE hFile = CreateFile(lpszFileName, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, 0, 0);
			if(hFile != INVALID_HANDLE_VALUE) {
				DWORD m_nFileSize = GetFileSize(hFile, 0);
				CloseHandle(hFile);
				if(m_pInput = new ifstream(lpszFileName, ios::in | ios::nocreate | ios::binary)) {
					if(m_pInput->fail()) {
						delete m_pInput;
						m_pInput = 0;
					} else
						hr = S_OK;
				}
			}
			delete[] lpszFileName;
		}
	} else
		hr = E_INVALIDARG;
	return hr;
}
*/
